"""
仿射模型使用线性回归(linear regression)
本例用简单的仿射模型进行数据推理,这里完全使用pytorch进行模型构造、训练
数据矩阵中,每行作为1个样本,列表示1个样本具有的feature
f(x)=x*w+b
loss=(f(xn)-y)**2
"""

import torch

# x/y数据为3*1的矩阵,3行1列
x_data = torch.Tensor([[1.0], [2.0], [3.0]])
y_data = torch.Tensor([[2.0], [4.0], [6.0]])


# 自定义线性回归模型,所有模型都需要继承自torch.nn.Module
class LinearModel(torch.nn.Module):
    # 必须要有的初始化
    def __init__(self, *args, **kwargs) -> None:
        # 必须先进行父类初始化
        super(LinearModel, self).__init__()
        # 使用torch已有的api进行线性回归模型定义
        # 这里的in/out对应x/y的列数(特征数)
        self.linear = torch.nn.Linear(1, 1)

    def forward(self, x):
        """
        前馈,这里会在epoch中model(x_data)计算时进行调用
        :param x: 输入的x
        :return: y_hat
        """
        y_hat = self.linear(x)
        return y_hat


# 初始化定义的线性回归模型
model = LinearModel()

# 定义损失函数,reduction="sum"时求取损失和,reduction="mean"时求取损失和的均值
criterion = torch.nn.MSELoss(reduction="sum")

# 定义优化器,model.parameters()会将所有模型下的权重找到,lr是学习率
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 开始训练
for epoch in range(100):
    # 使用线性回归模型,输入整个x数据集,计算出所有y_hat数据集
    y_pred = model(x_data)
    # 计算整个y_hat数据集损失
    loss = criterion(y_pred, y_data)
    # 输出当前的训练轮数和损失
    print(epoch, loss.item())
    # 将上一次的梯度归零
    optimizer.zero_grad()
    # 计算反向传播
    loss.backward()
    # 更新梯度
    optimizer.step()

# 输出当前的w/b
print("w: ", model.linear.weight.item())
print("b: ", model.linear.bias.item())

x_test = torch.Tensor([4.0])
y_test = model(x_test)
print("y_pred 4:", y_test.data)
